from collections import Iterator, Callable, Iterable
from numbers import Integral
from collections import deque
from multiprocessing import Pool


class preloaded_map(Iterator):
    def __init__(self, func, *iterables, preloaded_len, **pool_kwargs):
        assert isinstance(func, Callable)
        assert all(isinstance(it, Iterable) for it in iterables)
        assert isinstance(preloaded_len, Integral) and preloaded_len > 0

        self._func = func
        self._iterator = zip(*iterables)

        self._queue = deque()
        self._pool = Pool(**pool_kwargs)

        for _ in range(preloaded_len):
            self._enqueue()

    def _enqueue(self):
        try:
            params = next(self._iterator)
        except StopIteration:
            pass
        else:
            res = self._pool.apply_async(self._func, args=params)
            self._queue.append(res)

    def _dequeue(self):
        if len(self._queue) == 0:
            raise StopIteration
        else:
            res = self._queue.popleft()
            return res.get()

    def __next__(self):
        try:
            item = self._dequeue()
        except StopIteration:
            raise
        else:
            self._enqueue()
            return item


if __name__ == '__main__':
    import time
    import itertools

    def f(ii):
        time.sleep(0.1)
        return ii

    gen = preloaded_map(f, itertools.count(0), preloaded_len=10, processes=4)
    s0 = time.time()
    time.sleep(0.1)
    print(time.time() - s0)
    for _, i in zip(range(15), gen):
        time.sleep(0.03)
        print(i, time.time() - s0)
